from contextlib import asynccontextmanager
from fastapi import FastAPI
from fastapi.responses import JSONResponse
from s2st import router as s2st_router
from s2tt import router as s2tt_router
from t2st import router as t2st_router
from t2tt import router as t2tt_router
from asr import router as asr_router
from model import init_model, model, processor, get_model_status
import uvicorn

@asynccontextmanager
async def lifespan(app: FastAPI):
    # Startup: Load the model and processor
    init_model()
    yield
    # Shutdown: Cleanup (optional, e.g., clear model from memory)
    global model, processor
    model = None
    processor = None

app = FastAPI(
    title="SeamlessM4T v2 API",
    description="API for multilingual and multimodal translation using SeamlessM4T v2",
    version="1.0.0",
    lifespan=lifespan
)

# Include routers for each task
app.include_router(s2st_router, prefix="/s2st", tags=["Speech-to-Speech"])
app.include_router(s2tt_router, prefix="/s2tt", tags=["Speech-to-Text"])
app.include_router(t2st_router, prefix="/t2st", tags=["Text-to-Speech"])
app.include_router(t2tt_router, prefix="/t2tt", tags=["Text-to-Text"])
app.include_router(asr_router, prefix="/asr", tags=["Automatic Speech Recognition"])

@app.get("/")
async def root():
    return {"message": "Welcome to SeamlessM4T v2 API. Available endpoints: /s2st, /s2tt, /t2st, /t2tt, /asr"}

@app.get("/health")
async def health_check():
    """Health check endpoint to verify model status"""
    status = get_model_status()
    return {
        "status": "healthy" if status["model_loaded"] and status["processor_loaded"] else "unhealthy",
        "model_status": status,
        "message": "Model and processor are ready" if status["model_loaded"] and status["processor_loaded"] else "Model or processor not loaded"
    }

def main():
    # Run the FastAPI application with uvicorn using import string
    uvicorn.run(
        "main:app",
        host="127.0.0.1",
        port=8000,
        reload=True  # Enable auto-reload for development
    )

if __name__ == "__main__":
    main()